{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "%matplotlib inline\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.nn import functional as  F\n",
    "from torch import optim\n",
    "import numpy as np \n",
    "from matplotlib import pyplot as plt \n",
    "import matplotlib.animation\n",
    "import math, random\n",
    "\n",
    "TIME_STEP = 10\n",
    "INPUT_SIZE = 1\n",
    "DEVICE = torch.device(\"cuda\" if  torch.cuda.is_available() else \"cpu\")\n",
    "H_SIZE = 64\n",
    "EPOCHS = 300\n",
    "h_state = None\n",
    "steps = np.linspace(0, np.pi*2, 256, dtype=np.float32)\n",
    "x_np = np.sin(steps)\n",
    "y_np = np.cos(steps)\n",
    "\n",
    "plt.figure(1)\n",
    "plt.suptitle('Sin and Cos ', fontsize = '18')\n",
    "plt.plot(steps, y_np, 'r-', label = 'target cos')\n",
    "plt.plot(steps, x_np, 'b-', label='input (sin')\n",
    "plt.legend(loc='best')\n",
    "plt.show()"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "class RNN(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(RNN, self).__init__()\n",
    "        self.rnn = nn.RNN(\n",
    "            input_size = INPUT_SIZE, \n",
    "            hidden_size = H_SIZE, \n",
    "            num_layers = 1,\n",
    "            batch_first = True, \n",
    "        )\n",
    "        self.out = nn.Linear(H_SIZE, 1)\n",
    "    def forward(self, x, h_state):\n",
    "            # x (batch, time_step, input_size)\n",
    "         # h_state (n_layers, batch, hidden_size)\n",
    "         # r_out (batch, time_step, hidden_size)\n",
    "        r_out, h_state = self.rnn(x,h_state)\n",
    "        outs = []\n",
    "        print(\"size 1 : \", r_out.size(1))\n",
    "        for time_step in range(r_out.size(1)):\n",
    "            outs.append(self.out(r_out[:, time_step, :]))\n",
    "        return torch.stack(outs, dim=1), h_state\n",
    "        # 也可使用以下这样的返回值\n",
    "         # r_out = r_out.view(-1, 32)\n",
    "         # outs = self.out(r_out)\n",
    "         # return outs, h_state"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "rnn = RNN().to(DEVICE)\n",
    "optimizer = torch.optim.Adam(rnn.parameters())\n",
    "criterion = nn.MSELoss()"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "print(\"device: \", DEVICE)\n",
    "rnn.train()\n",
    "plt.figure(2)\n",
    "for step in range(EPOCHS):\n",
    "    start,end = step*np.pi, (step+1)*np.pi\n",
    "    steps = np.linspace(start, end, TIME_STEP, dtype=np.float32)\n",
    "    x_np = np.sin(steps)\n",
    "    y_np = np.cos(steps)\n",
    "    x = torch.from_numpy(x_np[np.newaxis, :, np.newaxis])  # shape (batch, time_step, input_size)\n",
    "    y = torch.from_numpy(y_np[np.newaxis, :, np.newaxis])\n",
    "    print(\"x size\", x.size())\n",
    "    x = x.to(DEVICE)\n",
    "    print(\"h_state first \", h_state)\n",
    "    prediction, h_state = rnn(x, h_state)  #h_state 在哪里定义？\n",
    "    print(\"h_state \", h_state, \"h_state data \", h_state.data)  #h_state.data 是没有之前h_state的梯度的\n",
    "    h_state = h_state.data # 重置隐藏层的状态, 切断和前一次迭代的链接\n",
    "    loss = criterion(prediction.cpu(), y)\n",
    "    optimizer.zero_grad()\n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "    if (step+1)%20 == 0:\n",
    "        print(\"EPOCHS: {}, Loss:{:4f}\".format(step,loss))\n",
    "        plt.plot(steps, y_np.flatten(), 'r-')\n",
    "        plt.plot(steps, prediction.cpu().data.numpy().flatten(),'b-')\n",
    "        plt.draw()\n",
    "        plt.pause(0.01)"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [],
   "outputs": [],
   "metadata": {}
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "0bedf3452ed48e38a2b75ff9f66ce8ae9ae2e92e8665cb618ee0797d2f46b9e5"
  },
  "kernelspec": {
   "name": "python3",
   "display_name": "Python 3.8.10 64-bit ('d2l-zh': conda)"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}